use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::TypeEnvironment;
use super::super::TypeInference;

impl TypeInference {



    /// Infer types in an if expression
    pub(in crate::type_check::type_inference) fn infer_if_expression(&mut self, if_expr: &IfExpr, env: &mut TypeEnvironment) -> Result<Type> {
        // Infer condition type
        let condition_type = self.infer_expression(&if_expr.condition, env)?;
        self.add_constraint(
            &condition_type,
            &Type::Basic(crate::ast::types::BasicType::Bool),
            if_expr.location,
        )?;

        // Infer then branch type
        let then_type = self.infer_expression(&if_expr.then_branch, env)?;

        // Infer else branch type
        if let Some(else_branch) = &if_expr.else_branch {
            let else_type = self.infer_expression(else_branch, env)?;
            self.add_constraint(&then_type, &else_type, if_expr.location)?;
        }

        Ok(then_type)
    }



    /// Infer types in a match expression
    pub(in crate::type_check::type_inference) fn infer_match_expression(
        &mut self,
        match_expr: &MatchExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        // Infer expression type
        let _expr_type = self.infer_expression(&match_expr.expr, env)?;

        if match_expr.arms.is_empty() {
            return Err(CompilerError::type_error(
                match_expr.location.line,
                match_expr.location.column,
                "Match expression must have at least one arm",
            ));
        }

        // Infer type of first arm
        let first_arm_type = self.infer_expression(&match_expr.arms[0].body, env)?;

        // Check that all arms have the same type
        for arm in &match_expr.arms[1..] {
            let arm_type = self.infer_expression(&arm.body, env)?;
            self.add_constraint(&arm_type, &first_arm_type, match_expr.location)?;
        }

        Ok(first_arm_type)
    }



    /// Infer types in a lambda expression
    pub(in crate::type_check::type_inference) fn infer_lambda_expression(
        &mut self,
        lambda: &LambdaExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        // Create new environment for lambda
        let mut lambda_env = env.child();

        // Add parameters to environment
        for param in &lambda.parameters {
            lambda_env.bind_variable(param.name.clone(), param.param_type.clone());
        }

        // Infer body type with optional try context based on lambda return type
        let ctx_depth = self.try_context_stack.len();
        self.push_try_context_from_type(lambda.return_type.as_ref());
        let body_type = {
            let result = self.infer_expression(&lambda.body, &mut lambda_env);
            self.truncate_try_context(ctx_depth);
            result?
        };

        // Create function type
        let param_types: Vec<Type> = lambda
            .parameters
            .iter()
            .map(|p| p.param_type.clone())
            .collect();
        let return_type = lambda.return_type.clone().unwrap_or(body_type);

        Ok(Type::Function(FunctionType {
            parameter_types: param_types,
            return_type: Some(Box::new(return_type)),
            variadic: false,
            location: lambda.location,
        }))
    }



    /// Infer types in a path expression (e.g., Result::Ok(value), Result::Err(error))
    pub(in crate::type_check::type_inference) fn infer_path_expression(
        &mut self,
        path_expr: &PathExpr,
        _env: &mut TypeEnvironment,
        expected_type: Option<&Type>,
    ) -> Result<Type> {
        // Path expressions like Result::Ok or Result::Err are handled in Call expressions
        // This function is called when the path itself is evaluated (without arguments)

        // Special handling for Option::None (which doesn't take arguments)
        if path_expr.segments.len() == 2 {
            if path_expr.segments[0] == "Option" && path_expr.segments[1] == "None" {
                if let Some(Type::Option(option_type)) = expected_type {
                    // Use contextual Option<T> so downstream consumers (e.g. `?`) know the exact type
                    return Ok(Type::Option(option_type.clone()));
                }
                // For None without context, return a generic marker so diagnostics can surface later
                return Ok(Type::Generic("__Option_None__".to_string()));
            }
        }

        // However, if this is just a type reference (not used as a constructor),
        // we might need to handle it differently
        // For now, return an error indicating that path expressions need context
        Err(CompilerError::type_error(
            path_expr.location.line,
            path_expr.location.column,
            format!(
                "Cannot infer type from path expression '{:?}' without context",
                path_expr.segments
            ),
        ))
    }

}
